// Copyright 2021-2024, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
//
// Redistribution and use in source and binary forms, with or without
// modification, are permitted provided that the following conditions
// are met:
//  * Redistributions of source code must retain the above copyright
//    notice, this list of conditions and the following disclaimer.
//  * Redistributions in binary form must reproduce the above copyright
//    notice, this list of conditions and the following disclaimer in the
//    documentation and/or other materials provided with the distribution.
//  * Neither the name of NVIDIA CORPORATION nor the names of its
//    contributors may be used to endorse or promote products derived
//    from this software without specific prior written permission.
//
// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY
// EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
// IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR
// PURPOSE ARE DISCLAIMED.  IN NO EVENT SHALL THE COPYRIGHT OWNER OR
// CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
// EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
// PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR
// PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY
// OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
// (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
// OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

#pragma once

#include <NvInfer.h>

#include <string>
#include <vector>

#include "shape_tensor.h"
#include "triton/backend/backend_common.h"
#include "triton/core/tritonserver.h"

namespace triton { namespace backend { namespace tensorrt {

TRITONSERVER_Error* GetProfileIndex(
    const std::string& profile_name, int* profile_index);

TRITONSERVER_DataType ConvertTrtTypeToDataType(nvinfer1::DataType trt_type);

std::string ConvertTrtTypeToConfigDataType(nvinfer1::DataType trt_type);

std::pair<bool, nvinfer1::DataType> ConvertDataTypeToTrtType(
    const TRITONSERVER_DataType& dtype);

bool CompareDims(
    const nvinfer1::Dims& model_dims, const std::vector<int64_t>& dims);
bool CompareDims(const nvinfer1::Dims& ldims, const nvinfer1::Dims& rdims);

TRITONSERVER_Error* ValidateDimension(
    const nvinfer1::Dims& this_dims, const nvinfer1::Dims& min_dims,
    const nvinfer1::Dims& max_dims);

// 'skip_first_dimension' should be set to true if the model has implicit batch
// dimension as 'this_dims' should be the full shape (contains batch dimension).
template <typename T>
TRITONSERVER_Error* ValidateDimension(
    const T& this_dims, const nvinfer1::Dims& min_dims,
    const nvinfer1::Dims& max_dims, const bool skip_first_dimension);

TRITONSERVER_Error* CompareDimsSupported(
    const std::string& model_name, const std::string& tensor_name,
    const nvinfer1::Dims& model_dims, common::TritonJson::Value& dims,
    const bool supports_batching, const bool compare_exact);

TRITONSERVER_Error* CompareDimsSupported(
    const std::string& model_name, const std::string& tensor_name,
    const nvinfer1::Dims& model_dims, const std::vector<int64_t>& dims,
    const bool supports_batching, const bool compare_exact);

TRITONSERVER_Error* CompareShapeDimsSupported(
    const std::string& model_name, const std::string& tensor_name,
    const nvinfer1::Dims& model_dims, common::TritonJson::Value& dims,
    const bool supports_batching);

TRITONSERVER_Error* ValidateControlDimsDynamic(
    const nvinfer1::Dims& dims, const bool support_batching);

TRITONSERVER_Error* ValidateShapeValues(
    const ShapeTensor& request_shape_values,
    const ShapeTensor& min_shape_values, const ShapeTensor& max_shape_values,
    size_t nb_shape_values);

template <typename T>
TRITONSERVER_Error* CheckShapeTensorInRange(
    const ShapeTensor& request_shape_values,
    const ShapeTensor& min_shape_values, const ShapeTensor& max_shape_values,
    size_t nb_shape_values);

void DimsToDimVec(const nvinfer1::Dims& model_dims, std::vector<int64_t>* dims);

TRITONSERVER_Error* DimsJsonToDimVec(
    common::TritonJson::Value& dims_json, std::vector<int64_t>* dims);

bool DimVecToDims(const std::vector<int64_t>& dim_vec, nvinfer1::Dims* dims);

int64_t GetElementCount(const nvinfer1::Dims& dims);

bool ContainsWildcard(const nvinfer1::Dims& dims);

bool ContainsWildcardAtExplicitBatchDim(const nvinfer1::Dims& dims);

const std::string DimsDebugString(const nvinfer1::Dims& dims);

const std::string DimsJsonToString(common::TritonJson::Value& dims);

const std::string TensorFormatToString(const nvinfer1::TensorFormat& io_format);

TRITONSERVER_Error* SupportsIntegratedZeroCopy(
    const int gpu_id, bool* zero_copy_support);

bool IsInput(nvinfer1::ICudaEngine* engine, const std::string& tensor_name);

bool ValidateBatchSize(
    const size_t total_batch_size, const ShapeTensor& min_shape_values,
    const ShapeTensor& max_shape_values);

//
// Templates
//

template <typename T>
TRITONSERVER_Error*
ValidateDimension(
    const T& this_dims, const nvinfer1::Dims& min_dims,
    const nvinfer1::Dims& max_dims, const bool skip_first_dimension)
{
  const int nonbatch_start_idx = (skip_first_dimension ? 1 : 0);
  if (int(this_dims.size()) != (max_dims.nbDims + nonbatch_start_idx)) {
    return TRITONSERVER_ErrorNew(
        TRITONSERVER_ERROR_INVALID_ARG,
        (std::string("model expected ") +
         std::to_string(max_dims.nbDims + nonbatch_start_idx) +
         " dimensions but received " + std::to_string(this_dims.size()) +
         " dimensions")
            .c_str());
  }

  for (int i = 0; i < max_dims.nbDims; i++) {
    if (this_dims[i] == -1) {
      continue;
    }
    if (this_dims[i + nonbatch_start_idx] < min_dims.d[i] ||
        this_dims[i + nonbatch_start_idx] > max_dims.d[i]) {
      return TRITONSERVER_ErrorNew(
          TRITONSERVER_ERROR_INVALID_ARG,
          (std::string("model expected the shape of dimension ") +
           std::to_string(i + nonbatch_start_idx) + " to be between " +
           std::to_string(min_dims.d[i]) + " and " +
           std::to_string(max_dims.d[i]) + " but received " +
           std::to_string(this_dims[i + nonbatch_start_idx]))
              .c_str());
    }
  }
  return nullptr;
}

template <typename T>
TRITONSERVER_Error*
CheckShapeTensorInRange(
    const ShapeTensor& request_shape_values,
    const ShapeTensor& min_shape_values, const ShapeTensor& max_shape_values,
    size_t nb_shape_values)
{
  const T* request_data =
      reinterpret_cast<const T*>(request_shape_values.GetData());
  const T* min_data = reinterpret_cast<const T*>(min_shape_values.GetData());
  const T* max_data = reinterpret_cast<const T*>(max_shape_values.GetData());
  for (size_t i = 0; i < nb_shape_values; i++) {
    if (request_data[i] < min_data[i] || request_data[i] > max_data[i]) {
      return TRITONSERVER_ErrorNew(
          TRITONSERVER_ERROR_INVALID_ARG,
          (std::string("The shape value at index ") + std::to_string(i) +
           " is expected to be in range from " + std::to_string(min_data[i]) +
           " to " + std::to_string(max_data[i]) +
           ", Got: " + std::to_string(request_data[i]))
              .c_str());
    }
  }
  return nullptr;
}

}}}  // namespace triton::backend::tensorrt
